import numpy as np
# import matplotlib.pyplot as plt
import time


# 响应谱计算函数 纯数组计算法
def resspec(Time, Amp, Tyz, Damp=0.05):  # 试着用纯数组方式处理
    """
     精确法求解响应谱（各个特征周期（频率）下的加速度响应最大值）
    :param Time: 时间序列
    :param Amp: 地震动时程幅值
    :param Tyz: 特征周期序列，即反应谱横坐标
    :param Damp:阻尼比   c/(2*sqrt(k*m))  代表结构在受激振后振动的衰减形式
    :return:
    """
    TA = Tyz.copy()
    (nt, na) = Amp.shape if Amp.ndim == 2 else (len(Amp), 0)
    nw = len(TA)
    if na:  # 同时输入了多条地震波  即Amp为多维
        # print('输入了多条地震波')
        # ***********精确法计算各反应***********
        # 初始化各储存向量 MDis为最大相对位移，MVel为最大相对速度，MAcc最大绝对加速度
        Displace = np.zeros((nt, nw, na))  # 相对位移    nt  nw   na   时间 频率 条数
        Velocity = np.zeros((nt, nw, na))  # 相对速度
        AbsAcce = np.zeros((nt, nw, na))  # 绝对加速度
        Dt = Time[1] - Time[0]  # 均匀时间间隔
        # 计算基本参数
        # 特征周期含有0，因此要先将其取出来，避免做分母
        TA = np.delete(TA, 0, axis=0)
        w = 2 * np.pi / TA  # 结构自振频率  从大到小排列
        # 再补一个位置回去
        TA = np.insert(TA, 0, 0, axis=0)
        w = np.insert(w, 0, 10000, axis=0)  # nw
        "=============下面开始计算参数矩阵==========="
        DamFrcy = w * np.sqrt(1 - Damp * Damp)  # 阻尼频率   nw
        e_t = np.exp(-Damp * w * Dt)  # nw
        s = np.sin(DamFrcy * Dt)  # nw
        c = np.cos(DamFrcy * Dt)  # nw
        A = np.zeros((2, 2, nw))  # 2，2，nw
        A[0, 0, :] = e_t * (s * Damp / np.sqrt(1 - Damp * Damp) + c)
        A[0, 1, :] = e_t * s / DamFrcy
        A[1, 0, :] = -w * e_t * s / np.sqrt(1 - Damp * Damp)
        A[1, 1, :] = e_t * (-s * Damp / np.sqrt(1 - Damp * Damp) + c)

        d_f = (2 * Damp ** 2 - 1) / (w ** 2 * Dt)
        d_3t = Damp / (w ** 3 * Dt)
        B = np.zeros((2, 2, nw))
        B[0, 0, :] = e_t * ((d_f + Damp / w) * s / DamFrcy + (2 * d_3t + 1 / w ** 2) * c) - 2 * d_3t
        B[0, 1, :] = -e_t * (d_f * s / DamFrcy + 2 * d_3t * c) - 1 / w ** 2 + 2 * d_3t
        B[1, 0, :] = e_t * ((d_f + Damp / w) * (c - Damp / np.sqrt(1 - Damp ** 2) * s) - (2 * d_3t + 1 / w ** 2) * (
                DamFrcy * s + Damp * w * c)) + 1 / (w ** 2 * Dt)
        B[1, 1, :] = e_t * (1 / (w ** 2 * Dt) * c + s * Damp / (w * DamFrcy * Dt)) - 1 / (w ** 2 * Dt)
        "===================== 根据地震记录, 计算反应======================"
        for i in range(nt - 1):  # 从1到nt-1 0另算
            Displace[i + 1, :, :] = A[0, 0, :].reshape(nw, 1) * Displace[i, :, :] + A[0, 1, :].reshape(nw, 1) \
                                    * Velocity[i, :, :] + B[0, 0, :].reshape(nw, 1) * Amp[i, :] \
                                    + B[0, 1, :].reshape(nw, 1) * Amp[i + 1, :]
            # 维度变化：  [nw  na]=[nw,1]*[nw,na] + [nw,1]*[1,nw,na] +[nw,1]*[1,na] +[nw,1]*[1,na]
            Velocity[i + 1, :, :] = A[1, 0, :].reshape(nw, 1) * Displace[i, :, :] + A[1, 1, :].reshape(nw, 1) \
                                    * Velocity[i, :, :] + B[1, 0, :].reshape(nw, 1) * Amp[i, :] + \
                                    B[1, 1, :].reshape(nw, 1) * Amp[i + 1, :]
            # 维度变化：  [1  nw  na]=[nw,1]*[1,nw,na] + [nw,1]*[1,nw,na] +[nw,1]*[1,na] +[nw,1]*[1,na]
            # 绝对加速度
            AbsAcce[i + 1, :, :] = -2 * Damp * w.reshape(nw, 1) * Velocity[i + 1, :, :] - \
                                   (w ** 2).reshape(nw, 1) * Displace[i + 1, :, :]
            # 维度变化：  [1  nw  na]=[nw,1]*[1,nw,na] - [nw,1]*[1,nw,na]
        # MDis = np.max(np.abs(Displace), axis=0)  # 求最大值
        # 反应谱中的点是对应各个特征周期的最大值
        # MVel = np.max(np.abs(Velocity), axis=0)
        MAcc = np.max(np.abs(AbsAcce), axis=0)
        # 反应谱第一点
        MAcc[0, :] = np.max(np.abs(Amp), axis=0)
        # 零周期加速度 ：指反应谱的高频未放大部分的加速度水平，它等于用于导出该反应谱所相应的时程曲线的最大峰值加速度。
    else:  # 只输入了一条地震波
        # print('只输入了一条地震波')
        # ***********精确法计算各反应***********
        # 初始化各储存向量
        Displace = np.zeros((nt, nw))  # 相对位移
        Velocity = np.zeros((nt, nw))  # 相对速度
        AbsAcce = np.zeros((nt, nw))  # 绝对加速度
        # 记录计算得到的反应，MDis为最大相对位移，MVel为最大相对速度，MAcc最大绝对加速度
        Dt = Time[1] - Time[0]

        # 特征周期含有0，因此要先将其取出来，避免做分母
        TA = np.delete(TA, 0, axis=0)
        w = 2 * np.pi / TA  # 结构自振频率  从高频到低频排列
        # 再补一个位置回去
        TA = np.insert(TA, 0, 0, axis=0)
        w = np.insert(w, 0, 10000, axis=0)

        DamFrcy = w * np.sqrt(1 - Damp * Damp)  # 阻尼频率
        e_t = np.exp(-Damp * w * Dt)
        s = np.sin(DamFrcy * Dt)
        c = np.cos(DamFrcy * Dt)
        # A矩阵
        A = np.zeros((2, 2, nw))
        A[0, 0, :] = e_t * (s * Damp / np.sqrt(1 - Damp * Damp) + c)
        A[0, 1, :] = e_t * s / DamFrcy
        A[1, 0, :] = -w * e_t * s / np.sqrt(1 - Damp * Damp)
        A[1, 1, :] = e_t * (-s * Damp / np.sqrt(1 - Damp * Damp) + c)
        # B矩阵
        d_f = (2 * Damp ** 2 - 1) / (w ** 2 * Dt)
        d_3t = Damp / (w ** 3 * Dt)
        B = np.zeros((2, 2, nw))
        B[0, 0, :] = e_t * ((d_f + Damp / w) * s / DamFrcy + (2 * d_3t + 1 / w ** 2) * c) - 2 * d_3t
        B[0, 1, :] = -e_t * (d_f * s / DamFrcy + 2 * d_3t * c) - 1 / w ** 2 + 2 * d_3t
        B[1, 0, :] = e_t * ((d_f + Damp / w) * (c - Damp / np.sqrt(1 - Damp ** 2) * s) - (2 * d_3t + 1 / w ** 2) * (
                DamFrcy * s + Damp * w * c)) + 1 / (w ** 2 * Dt)
        B[1, 1, :] = e_t * (1 / (w ** 2 * Dt) * c + s * Damp / (w * DamFrcy * Dt)) - 1 / (w ** 2 * Dt)

        for i in range(nt - 1):  # 根据地震记录, 计算反应
            Displace[i + 1, :] = A[0, 0, :] * Displace[i, :] + A[0, 1, :] * Velocity[i, :] \
                                 + B[0, 0, :] * Amp[i] + B[0, 1, :] * Amp[i + 1]
            Velocity[i + 1, :] = A[1, 0, :] * Displace[i, :] + A[1, 1, :] * Velocity[i, :] + \
                                 B[1, 0, :] * Amp[i] + B[1, 1, :] * Amp[i + 1]
            AbsAcce[i + 1, :] = -2 * Damp * w * Velocity[i + 1, :] - w ** 2 * Displace[i + 1, :]

        # MDis = np.max(np.abs(Displace), axis=0)  # 最大位移
        # 反应谱中的点是对应各个特征周期的最大值
        # MVel = np.max(np.abs(Velocity), axis=0)  # 最大速度
        MAcc = np.max(np.abs(AbsAcce), axis=0)
        MAcc[0] = np.max(np.abs(Amp))  # 零周期加速度  高频  直接取加速度时程的最大值

    return MAcc


# 响应谱计算函数 原方法
def resspec2(Time, Amp, TA, Damp=0.05):
    """
     精确法求解响应谱（各个特征周期（频率）下的加速度响应最大值）
    :param Time: 时间序列
    :param Amp: 地震动时程幅值
    :param TA: 特征周期序列，即反应谱横坐标
    :param Damp:阻尼比   c/(2*sqrt(k*m))  代表结构在受激振后振动的衰减形式
    :return:
    """
    count = len(Amp)
    nw = len(TA)
    # ***********精确法计算各反应***********
    # 初始化各储存向量
    Displace = np.zeros(count)  # 相对位移
    Velocity = np.zeros(count)  # 相对速度
    AbsAcce = np.zeros(count)  # 绝对加速度
    # 记录计算得到的反应，MDis为最大相对位移，MVel为最大相对速度，MAcc最大绝对加速度
    MDis = np.zeros(nw)
    MVel = np.zeros(nw)
    MAcc = np.zeros(nw)

    Dt = Time[1] - Time[0]
    t = 0
    for T in TA:  # 变量 不同的结构周期
        w = 2 * np.pi / T  # 结构自振频率
        DamFrcy = w * np.sqrt(1 - Damp * Damp)  # 阻尼频率
        e_t = np.exp(-Damp * w * Dt)
        s = np.sin(DamFrcy * Dt)
        c = np.cos(DamFrcy * Dt)
        A = np.zeros((2, 2))
        A[0, 0] = e_t * (s * Damp / np.sqrt(1 - Damp * Damp) + c)
        A[0, 1] = e_t * s / DamFrcy
        A[1, 0] = -w * e_t * s / np.sqrt(1 - Damp * Damp)
        A[1, 1] = e_t * (-s * Damp / np.sqrt(1 - Damp * Damp) + c)
        d_f = (2 * Damp ** 2 - 1) / (w ** 2 * Dt)
        d_3t = Damp / (w ** 3 * Dt)
        B = np.zeros((2, 2))
        B[0, 0] = e_t * ((d_f + Damp / w) * s / DamFrcy + (2 * d_3t + 1 / w ** 2) * c) - 2 * d_3t
        B[0, 1] = -e_t * (d_f * s / DamFrcy + 2 * d_3t * c) - 1 / w ** 2 + 2 * d_3t
        B[1, 0] = e_t * ((d_f + Damp / w) * (c - Damp / np.sqrt(1 - Damp ** 2) * s) - (2 * d_3t + 1 / w ** 2) * (
                DamFrcy * s + Damp * w * c)) + 1 / (w ** 2 * Dt)
        B[1, 1] = e_t * (1 / (w ** 2 * Dt) * c + s * Damp / (w * DamFrcy * Dt)) - 1 / (w ** 2 * Dt)
        for i in range(count - 1):  # 根据地震记录, 计算反应
            Displace[i + 1] = A[0, 0] * Displace[i] + A[0, 1] * Velocity[i] + B[0, 0] * Amp[i] + B[0, 1] * Amp[i + 1]
            Velocity[i + 1] = A[1, 0] * Displace[i] + A[1, 1] * Velocity[i] + B[1, 0] * Amp[i] + B[1, 1] * Amp[i + 1]
            AbsAcce[i + 1] = -2 * Damp * w * Velocity[i + 1] - w ** 2 * Displace[i + 1]

        MDis[t] = max(np.abs(Displace))  # 求最大值
        # 反应谱中的点是对应各个特征周期的最大值
        MVel[t] = max(np.abs(Velocity))
        if T == 0.0:  # 对初始值设置
            MAcc[t] = max(np.abs(Amp))
        else:
            MAcc[t] = max(np.abs(AbsAcce))
        Displace = np.zeros(count)  # 初始化各储存向量，避免下次不同周期计算时引用到前一个周期的结果
        # 各个周期的结果已经储存到MDis和MVel中了
        Velocity = np.zeros(count)
        AbsAcce = np.zeros(count)
        t += 1
    return MAcc, MVel, MDis


# # 测试
# 导入数据
if __name__ == '__main__':
    earthdata = np.loadtxt('..\实际地震波\el-centro.txt')  # elcentro波  单位mm/s^2
    # earthdata = np.loadtxt('C:/Users/francklinson/PycharmProjects/pytorch-gpu/地震波生成/ninghedizhen.txt')  # 宁河波，单位cm/s^2
    # earthdata = np.loadtxt('C:/Users/Administrator/Desktop/data.txt')
    # print(len(earthdata))
    t = earthdata[:2000, 0]
    n = 50  # na
    amp = np.zeros((2000, n))
    for i in range(n):
        amp[:, i] = earthdata[:2000, 1]
    ampx = earthdata[:2000, 1]
    ta = np.arange(0.02, 6.02, 0.02)
    # print(len(ta))  # nw
    t1 = time.time()
    for i in range(n):
        resspec2(t, amp[:, i], ta)
    t2 = time.time()
    resspec(t, amp, ta)
    t3 = time.time()

    print(t2 - t1)
    print(t3 - t2)
    print((t2 - t1) / (t3 - t2))
